#include <iostream>
#include "util/distance.hpp"
#include "util/time_memory.hpp"
#include <vector>
#include <cstring>
#include <cmath>
#include <unordered_set>
#include <queue>
#include <random>

#undef NDEBUG

#include <cassert>
#include <algorithm>

/*
 * 假设当前迭代点为cand, 邻居为n, 查询为q
 * 终止条件:
 *      向量cand q在cand n上的投影与向量cand n的模进行比较, 如果小就全部搜索, 否则只搜索梯度边
 */

namespace GraphGradient {

    class GraphBaseline {
    public:
        GraphBaseline(int vec_dim, int max_element, int random_seed = 100) {
            max_element_ = max_element;
            vec_dim_ = vec_dim;
            srand(random_seed);
            report_every_ = 100000;
        }

        GraphBaseline() {

        }

        void add(float *data, std::vector<std::vector<int>> &graph, int enter_point);

        std::vector<std::unordered_set<int>> search_knn(float *query, int k, int n_query, int ef_search);

        ~GraphBaseline() {
        }

        std::vector<std::vector<int>> graph_;

        unsigned int test_n_cand_ = 0;
        unsigned int test_n_neighbor_ = 0;
#ifdef TEST
        bool is_warm_up_ = true;
        unsigned int test_avg_n_hops_ = 0;
        unsigned int test_avg_n_visit_edges_ = 0;
        unsigned int test_avg_n_visited_edges_ = 0;
        std::vector<unsigned int> test_avg_n_hops_l;
        std::vector<unsigned int> test_avg_n_visit_edges_l;
        std::vector<unsigned int> test_avg_n_visited_edges_l;
#endif

    private:
        float *data_;
        int max_element_;
        int vec_dim_;
        int enterpoint_;
        size_t report_every_;

        std::unordered_set<int> search_labeled(float *item, int k, int ef_search);

    };

    /*
     * use the distance to enter point as the judgement of gradient and neighbor
     */
    void GraphBaseline::add(float *data, std::vector<std::vector<int>> &graph, int enter_point) {
        data_ = data;
        enterpoint_ = enter_point;

        float *ep_data = data_ + enterpoint_ * vec_dim_;
        std::vector<int> distance2ep(max_element_);
        for (int i = 0; i < max_element_; i++) {
            distance2ep[i] = l2distance(data + i * vec_dim_, ep_data, vec_dim_);
        }

        graph_.resize(max_element_);
        TimeRecord single_insert_record;
        for (int i = 0; i < max_element_; i++) {
            std::vector<int> graph_i = graph[i];
            int graph_i_size = graph_i.size();
            graph_[i].resize(graph_i_size);
            int start_idx = 0;
            for (std::vector<int>::iterator it = graph_i.begin(); it != graph_i.end(); ++it) {
                graph_[i][start_idx] = *it;
                start_idx++;
            }

            if (i % report_every_ == 0) {
                std::cout << i / (0.01 * max_element_) << " %, "
                          << 1e-6 * single_insert_record.get_elapsed_time_micro() << " s/iter" << " Mem: "
                          << get_current_RSS() / 1000000 << " Mb \n";
                single_insert_record.reset();
            }
        }

        for (int i = 0; i < max_element_; i++) {
            graph[i].clear();
        }
        graph.clear();

    }

    std::vector<std::unordered_set<int>> GraphBaseline::search_knn(float *query, int k, int n_query, int ef_search) {
        std::vector<std::unordered_set<int>> res(n_query);
        float *tmp_query = query;
        for (int i = 0; i < n_query; i++, tmp_query += vec_dim_) {
            res[i] = search_labeled(tmp_query, k, ef_search);
        }
#ifdef TEST
        test_avg_n_hops_ = test_avg_n_hops_ / n_query;
        test_avg_n_visit_edges_ = test_avg_n_visit_edges_ / n_query;
        test_avg_n_visited_edges_ = test_avg_n_visited_edges_ / n_query;
        if(!is_warm_up_){
            test_avg_n_hops_l.push_back(test_avg_n_hops_);
            test_avg_n_visit_edges_l.push_back(test_avg_n_visit_edges_);
            test_avg_n_visited_edges_l.push_back(test_avg_n_visited_edges_);
        }
        test_avg_n_hops_ = 0;
        test_avg_n_visit_edges_ = 0;
        test_avg_n_visited_edges_ = 0;
#endif

        return res;
    }

    //返回的数组是按照相对query距离排序的, 降序, 就是第一个元素是数组中离query最远的元素
    std::unordered_set<int> GraphBaseline::search_labeled(float *item, int k, int ef_search) {
        std::unordered_set<int> visited_element;
        std::priority_queue<distance_cmp, std::vector<distance_cmp>, distance_min_heap> candidates;
        std::priority_queue<distance_cmp, std::vector<distance_cmp>, distance_max_heap> nearest_neighbors;
        double dist_ep_query = l2distance(data_ + enterpoint_ * vec_dim_, item, vec_dim_);
        {
            visited_element.insert(enterpoint_);
            candidates.emplace(dist_ep_query, enterpoint_);
            nearest_neighbors.emplace(dist_ep_query, enterpoint_);
        }
        while (candidates.size() > 0) {
            distance_cmp cand = candidates.top();
#ifdef TEST
            test_avg_n_hops_++;
#endif
            candidates.pop();
            distance_cmp far_ele = nearest_neighbors.top();
            if (cand.distance_ > far_ele.distance_) {
                break;
            }
            std::vector<int> neighbor_l = graph_[cand.idx_];
            int end_idx = neighbor_l.size();
            for (int i = 0; i < end_idx; i++) {
                int neighbor = neighbor_l[i];
                if (visited_element.find(neighbor) == visited_element.end()) { //can not find
#ifdef TEST
                    test_avg_n_visit_edges_++;
#endif
                    visited_element.insert(neighbor);
                    far_ele = nearest_neighbors.top();
                    double neighbor_dist = l2distance(item, data_ + neighbor * vec_dim_, vec_dim_);
                    if (neighbor_dist < far_ele.distance_ || (int) nearest_neighbors.size() < ef_search) {

                        candidates.emplace(neighbor_dist, neighbor);
                        nearest_neighbors.emplace(neighbor_dist, neighbor);
                        while ((int) nearest_neighbors.size() > ef_search) {
                            nearest_neighbors.pop();
                        }
                    }
                }
#ifdef TEST
                test_avg_n_visited_edges_++;
#endif
            }

        }
        while ((int) nearest_neighbors.size() > k) {
            nearest_neighbors.pop();
        }

        std::unordered_set<int> res;
        while (nearest_neighbors.size() > 0) {
            distance_cmp tmp = nearest_neighbors.top();
            nearest_neighbors.pop();
            res.insert(tmp.idx_);
        }
        return res;
    }

}